import os
import numpy as np
import nibabel as nib
import cv2
from pathlib import Path

# === 路径配置 ===
# 原始PNG路径
image_dir = Path("/data/suziren25/suziren/nnUNet/nnUNet_raw/abdominal_US/new_images")
# 推理结果（mask）路径
mask_dir = Path("/data/suziren25/suziren/nnUNet/nnUNet_results/Dataset506_KidneyUS/inference_2d_PP")
# 输出路径
output_dir = Path("/data/suziren25/suziren/nnUNet/visualization/inference_overlay_new")
output_dir.mkdir(parents=True, exist_ok=True)

# === 可视化参数 ===
overlay_color = (255, 0, 0)  # 红色
alpha = 0.5  # 透明度

for mask_path in sorted(mask_dir.glob("*.nii.gz")):
    # ⚙️ 去掉双扩展名（.nii.gz）
    name = mask_path.name.replace(".nii.gz", "")
    img_name = f"{name}.png"
    img_path = image_dir / img_name

    if not img_path.exists():
        print(f"⚠️ 原图不存在: {img_path}")
        continue

    # === 读取mask ===
    mask_nifti = nib.load(str(mask_path))
    mask_data = mask_nifti.get_fdata()
    mask_data = (mask_data > 0).astype(np.uint8)

    # === 读取原图 ===
    img = cv2.imread(str(img_path))
    if img is None:
        print(f"⚠️ 无法读取图片: {img_path}")
        continue

    # === 调整尺寸一致 ===
    mask_resized = cv2.resize(mask_data, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_NEAREST)

    # === 生成叠加图 ===
    overlay = img.copy()
    overlay[mask_resized == 1] = overlay_color
    blended = cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0)

    # === 保存 ===
    out_path = output_dir / img_name
    cv2.imwrite(str(out_path), blended)
    print(f"✅ 已生成可视化结果: {out_path}")

print("\n🎉 全部完成！请查看输出目录：", output_dir)